#!/usr/bin/python
# -*- encoding:utf-8 -*-
'''
Author: Alexxrhuang
Data: 2019.07.2
Updata:
    采用audioop.ratecv进行下采样
Note:
    该算法得到帧级别的标注，然后一块一块的输出非静音部分
    比较难得到准确的段时间，但是有帧级别的标注
'''


from __future__ import print_function
import sys
import wave
import audioop
import webrtcvad
import collections
import contextlib


def read_wave(path):
    """Reads a .wav file.

    Takes the path, and returns (PCM audio data, sample rate).
    """
    sr_lists = [8000, 16000, 32000, 48000]
    with contextlib.closing(wave.open(path, 'rb')) as wf:
        num_channels = wf.getnchannels()
        assert num_channels == 1
        sample_width = wf.getsampwidth()
        assert sample_width == 2
        sr = wf.getframerate()
        pcm_data = wf.readframes(wf.getnframes())           # pcm_data: bytes array
        target_sr = sr
        if sr not in sr_lists:
            if sr < sr_lists[0]:
                print('Error: sample rate less than 8000')
                exit(0)
            for i in sr_lists:
                if i < sr:
                    target_sr = i
            pcm_data = audioop.ratecv(pcm_data, sample_width, num_channels, sr, target_sr, None)
            if num_channels == 1:
                pcm_data = audioop.tomono(pcm_data[0], sample_width, num_channels, 0)
            print('Info: resample from %d to %d' % (sr, target_sr))

    return pcm_data, target_sr


def write_wave(path, audio, sample_rate):
    """Writes a .wav file.

    Takes path, PCM audio data, and sample rate.
    """
    with contextlib.closing(wave.open(path, 'wb')) as wf:
        wf.setnchannels(1)
        wf.setsampwidth(2)
        wf.setframerate(sample_rate)
        wf.writeframes(audio)


class Frame(object):
    """Represents a "frame" of audio data."""
    def __init__(self, bytes, timestamp, duration):
        self.bytes = bytes
        self.timestamp = timestamp
        self.duration = duration


def frame_generator(frame_duration_ms, audio, sample_rate):
    """Generates audio frames from PCM audio data.

    Takes the desired frame duration in milliseconds, the PCM data, and
    the sample rate.

    Yields Frames of the requested duration.
    """
    n = int(sample_rate * (frame_duration_ms / 1000.0) * 2)     # 每一帧的样本点数
    offset = 0
    timestamp = 0.0
    duration = (float(n) / sample_rate) / 2.0                   # 每一帧的时间
    while offset + n < len(audio):
        yield Frame(audio[offset:offset + n], timestamp, duration)
        timestamp += duration
        offset += n


def vad_collector(sample_rate, frame_duration_ms,
                  padding_duration_ms, vad, frames):
    """Filters out non-voiced audio frames.

    Given a webrtcvad.Vad and a source of audio frames, yields only
    the voiced audio.

    Uses a padded, sliding window algorithm over the audio frames.
    When more than 90% of the frames in the window are voiced (as
    reported by the VAD), the collector triggers and begins yielding
    audio frames. Then the collector waits until 90% of the frames in
    the window are unvoiced to detrigger.

    The window is padded at the front and back to provide a small
    amount of silence or the beginnings/endings of speech around the
    voiced frames.

    Arguments:

    sample_rate - The audio sample rate, in Hz.
    frame_duration_ms - The frame duration in milliseconds, 帧长，default 30ms, 
    padding_duration_ms - The amount to pad the window, in milliseconds, 滑动窗长度，default 300ms
    vad - An instance of webrtcvad.Vad.
    frames - a source of audio frames (sequence or generator).

    Returns: A generator that yields PCM audio data.
    """
    num_padding_frames = int(padding_duration_ms / frame_duration_ms)
    # We use a deque for our sliding window/ring buffer.
    ring_buffer = collections.deque(maxlen=num_padding_frames)
    # We have two states: TRIGGERED and NOTTRIGGERED. We start in the
    # NOTTRIGGERED state.
    triggered = False

    n = int(sample_rate * (frame_duration_ms / 1000.0))
    duration = 2 * float(n) / sample_rate
    voiced_frames = []
    frame_start = -1
    frame_end = -1
    frame_count = 0
    for frame in frames:
        frame_count += 1
        is_speech = vad.is_speech(frame.bytes, sample_rate)
        # sys.stdout.write('1' if is_speech else '0')
        # 如果还没有进入语音区
        if not triggered:
            ring_buffer.append((frame, is_speech))
            num_voiced = len([f for f, speech in ring_buffer if speech])        # 统计当前ring缓冲区的语音帧数目
            # If we're NOTTRIGGERED and more than 90% of the frames in
            # the ring buffer are voiced frames, then enter the
            # TRIGGERED state.
            # 触发语音区
            if num_voiced > 0.9 * ring_buffer.maxlen:
                triggered = True
                # sys.stdout.write('+(%s)' % (ring_buffer[0][0].timestamp,))
                # We want to yield all the audio we see from now until
                # we are NOTTRIGGERED, but we have to start with the
                # audio that's already in the ring buffer.
                # 将缓冲区之前存储的帧加入voiced_frames中
                for f, s in ring_buffer:
                    voiced_frames.append(f)
                ring_buffer.clear()
                frame_start = frame_count - len(voiced_frames)
        else:
            # We're in the TRIGGERED state, so collect the audio data
            # and add it to the ring buffer.
            # 语音区的帧默认加入voiced_frames中
            voiced_frames.append(frame)
            ring_buffer.append((frame, is_speech))
            num_unvoiced = len([f for f, speech in ring_buffer if not speech])  # 统计当前ring缓冲区的静音帧数目
            # If more than 90% of the frames in the ring buffer are
            # unvoiced, then enter NOTTRIGGERED and yield whatever
            # audio we've collected.
            if num_unvoiced > 0.9 * ring_buffer.maxlen:
                # sys.stdout.write('-(%s)' % (frame.timestamp + frame.duration))
                triggered = False
                yield b''.join([f.bytes for f in voiced_frames])                # 输出整个语音区所有的帧
                ring_buffer.clear()
                voiced_frames = []
                frame_end = frame_count + len(voiced_frames)
                print(frame_start*duration, frame_end*duration)
                frame_end = -1
                frame_start = -1
    # if triggered:
    #     sys.stdout.write('-(%s)' % (frame.timestamp + frame.duration))
    
    sys.stdout.write('\n')
    # If we have any leftover voiced audio when we run out of input,
    # yield it.
    # 如果在wav音频结束时语音区仍然有数据，即一直没有触发重新进行静音区
    if voiced_frames:
        print(frame_start*duration, frame_end*duration)
        yield b''.join([f.bytes for f in voiced_frames])


def main(args):
    if len(args) != 2:
        sys.stderr.write(
            'Usage: example.py <aggressiveness> <path to wav file>\n')
        sys.exit(1)
    audio, sample_rate = read_wave(args[1])
    vad = webrtcvad.Vad(int(args[0]))
    frames = frame_generator(30, audio, sample_rate)
    frames = list(frames)
    segments = vad_collector(sample_rate, 30, 300, vad, frames)
    for i, segment in enumerate(segments):
        path = 'chunk-%002d.wav' % (i,)
        print(' Writing %s' % (path,))
        write_wave(path, segment, sample_rate)


if __name__ == '__main__':
    main(sys.argv[1:])
